{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "根据关键词对企业分类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-05 14:32:12,358 - modelscope - INFO - PyTorch version 2.3.0 Found.\n",
      "2024-09-05 14:32:12,359 - modelscope - INFO - Loading ast index from /home/jie/.cache/modelscope/ast_indexer\n",
      "2024-09-05 14:32:12,375 - modelscope - INFO - Loading done! Current index file version is 1.14.0, with md5 9f53011443b7fabdfb807148bea6e1f2 and a total number of 976 components indexed\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from industries import hydrogen as industry\n",
    "from prompt import get_industry_trans_func\n",
    "from utils import glm4_vllm, load_obj\n",
    "from setting import StaticValues\n",
    "\n",
    "sv = StaticValues(industry.NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = sv.logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '英文名称', '统一社会信用代码', '企业类型', '经营状态', '成立日期', '核准日期', '法定代表人', '注册资本', '实缴资本', '参保人数', '公司规模', '经营范围', '注册地址', '营业期限', '纳税人识别号', '工商注册号', '组织机构代码', '纳税人资质', '曾用名', '所属省份', '所属城市', '所属区县', '网站链接', '所属行业', '一级行业分类', '二级行业分类', '三级行业分类', '登记机关', '经度', '纬度', 'is_industry_prompt'],\n",
       "    num_rows: 4115\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 导入关键词分类的表\n",
    "dataset = load_dataset(\"csv\", data_files=sv.KW_CSV, split=\"train\")\n",
    "new_dataset = dataset.map(\n",
    "    get_industry_trans_func(industry.binary_cls_prompt),\n",
    "    )\n",
    "new_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "industry.binary_cls_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = []\n",
    "for item in new_dataset:\n",
    "    prompts.append([{\"role\": \"user\", \"content\": item[\"is_industry_prompt\"]}])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 封装一下，如果已经数据，自动导入\n",
    "glm4_vllm(prompts, output_dir=sv.BINARY_VLLM_OBJ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "vllm_pred = load_obj(sv.BINARY_VLLM_OBJ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "reason:\n",
      "根据提供的企业信息，安新县政屿农业科技有限公司的主要业务涉及农业种植、技术服务、仓储服务、租赁服务、销售预包装食品等，与氢能企业的关键特征和业务范畴不符。氢能企业通常专注于氢气的生产、储存、运输、加注、使用和销售，而该公司的业务主要集中在农业领域。\n",
      "\n",
      "label: 否\n"
     ]
    }
   ],
   "source": [
    "print(vllm_pred[0].outputs[0].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def reason_label_parse(text):\n",
    "    # 定义正则表达式来匹配reason和label\n",
    "    pattern = r\"(reason:)(.*?)(?=\\nlabel:|$)\"\n",
    "    pattern_label = r\"(label:)(.*?)$\"\n",
    "\n",
    "    # 使用正则表达式查找reason\n",
    "    match_reason = re.search(pattern, text, re.DOTALL)\n",
    "    if match_reason:\n",
    "        reason = match_reason.group(2).strip()  # 去除前后的空白字符\n",
    "\n",
    "    # 使用正则表达式查找label\n",
    "    match_label = re.search(pattern_label, text, re.DOTALL)\n",
    "    if match_label:\n",
    "        label = match_label.group(2).strip()  # 去除前后的空白字符\n",
    "        chars_to_remove = [\"\\\"\", \"[\", \"]\", \"'\"]  \n",
    "        for char in chars_to_remove:  \n",
    "            label = label.replace(char, \"\")\n",
    "\n",
    "    return reason, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vllm_pred[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/05 14:44:14 - hydrogen - INFO: \n",
      "在通过关键词筛选出的4115个企业样本中，通过大模型分类，其中2个是hydrogen企业，4113个不是hydrogen企业。\n"
     ]
    }
   ],
   "source": [
    "yes_cnt = 0\n",
    "no_cnt = 0\n",
    "\n",
    "reasons = []\n",
    "labels = []\n",
    "\n",
    "idxs = []\n",
    "for idx, pred in enumerate(vllm_pred):\n",
    "    pred_text = pred.outputs[0].text\n",
    "    reason, label = reason_label_parse(pred_text)\n",
    "    reasons.append(reason)\n",
    "    labels.append(label)\n",
    "    if label == \"是\":\n",
    "        yes_cnt += 1\n",
    "        idxs.append(idx)\n",
    "    elif label == \"否\":\n",
    "        no_cnt += 1\n",
    "\n",
    "if yes_cnt + no_cnt == len(vllm_pred):\n",
    "    logger.info(\n",
    "        f\"在通过关键词筛选出的{len(vllm_pred)}个企业样本中，通过大模型分类，其中{yes_cnt}个是{sv.NAME}企业，{no_cnt}个不是{sv.NAME}企业。\"\n",
    "    )\n",
    "else:\n",
    "    logger.info(\n",
    "        f\"在通过关键词筛选出的{len(vllm_pred)}个企业样本中，通过大模型分类，其中{yes_cnt}个是{sv.NAME}企业，{no_cnt}个不是{sv.NAME}企业。还有{len(vllm_pred) - yes_cnt - no_cnt}个样本无法确定。\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2817, 3222]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_dataset = dataset.add_column(\"reason\", reasons)\n",
    "tmp_dataset = tmp_dataset.add_column(\"label\", labels)\n",
    "# 下标筛选\n",
    "tmp_dataset = tmp_dataset.select(idxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '英文名称', '统一社会信用代码', '企业类型', '经营状态', '成立日期', '核准日期', '法定代表人', '注册资本', '实缴资本', '参保人数', '公司规模', '经营范围', '注册地址', '营业期限', '纳税人识别号', '工商注册号', '组织机构代码', '纳税人资质', '曾用名', '所属省份', '所属城市', '所属区县', '网站链接', '所属行业', '一级行业分类', '二级行业分类', '三级行业分类', '登记机关', '经度', '纬度', 'reason', 'label'],\n",
       "    num_rows: 2\n",
       "})"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'根据提供的企业信息，雄安新动力科技股份有限公司的业务范围涵盖了多个领域，包括氢能源领域内技术的研发、推广和服务，以及节能技术研发、技术咨询、技术服务等。这些信息表明该企业确实在氢能源领域有所涉足，尤其是在氢能源技术的研发和服务方面。\\n\\n然而，企业信息中并未明确提及该企业具体涉及氢能的制氢技术、储运氢能力、氢能源电池技术、氢能基础设施的建设和运营，以及氢能产业生态的构建。虽然氢能源领域的技术研发和服务是氢能企业的重要组成部分，但仅凭这些信息无法确定该企业是否在氢能产业链的其它环节有所涉及。'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_dataset[0][\"reason\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7faefa9b8ddf43d89a2d25de722b2c93",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating CSV from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "4620"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_dataset.to_csv(\"tmp.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
